package org.easy.milvus;

import com.alibaba.fastjson.JSON;
import org.easy.core.exception.BizRuntimeException;
import org.easy.core.util.Lists;
import org.easy.core.util.ObjectUtil;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.DataType;
import io.milvus.grpc.MutationResult;
import io.milvus.grpc.SearchResults;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.RpcStatus;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.DropCollectionParam;
import io.milvus.param.collection.FieldType;
import io.milvus.param.collection.HasCollectionParam;
import io.milvus.param.collection.LoadCollectionParam;
import io.milvus.param.collection.ReleaseCollectionParam;
import io.milvus.param.dml.DeleteParam;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.param.highlevel.dml.GetIdsParam;
import io.milvus.param.highlevel.dml.response.GetResponse;
import io.milvus.param.index.CreateIndexParam;
import io.milvus.response.QueryResultsWrapper;
import lombok.Getter;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

@Getter
public class MilvusTemplate {

    private final MilvusServiceClient milvusServiceClient;

    private final EmbeddingClient embeddingClient;

    public String wrapCollectionName(String collectionName) {
        return "c_" + collectionName;
    }


    public MilvusTemplate(MilvusServiceClient client, EmbeddingClient embeddingClient) {
        this.milvusServiceClient = client;
        this.embeddingClient = embeddingClient;
    }


    public boolean hasCollection(String collectionName) {
        return milvusServiceClient.hasCollection(HasCollectionParam.newBuilder().withCollectionName(wrapCollectionName(collectionName)).build()).getData();
    }

    public void createDefaultCollection(String collectionName) {
        R<Boolean> booleanR = milvusServiceClient.hasCollection(HasCollectionParam.newBuilder().withCollectionName(wrapCollectionName(collectionName)).build());
        if (booleanR.getData()) {
            return;
        }

        CreateCollectionParam createCollectionReq = CreateCollectionParam.newBuilder()
                .withCollectionName(wrapCollectionName(collectionName))
                .addFieldType(FieldType.newBuilder()
                        .withName("id")
                        .withDataType(DataType.Int64)
                        .withPrimaryKey(true)
                        .withAutoID(true)
                        .build())
                .addFieldType(FieldType.newBuilder()
                        .withName("content")
                        .withDataType(DataType.VarChar)
                        .withMaxLength(512)
                        .build())
                .addFieldType(FieldType.newBuilder()
                        .withName("content_vector_question")
                        .withDataType(DataType.FloatVector)
                        .withDimension(768)
                        .build())
                .addFieldType(FieldType.newBuilder()
                        .withName("content_vector_answer")
                        .withDataType(DataType.FloatVector)
                        .withDimension(768)
                        .build())
                .addFieldType(FieldType.newBuilder()
                        .withName("partition_name")
                        .withDataType(DataType.VarChar)
                        .withPartitionKey(true)
                        .withMaxLength(256)
                        .build())
                .build();
        milvusServiceClient.createCollection(createCollectionReq);

        milvusServiceClient.createIndex(
                CreateIndexParam.newBuilder()
                        .withCollectionName(wrapCollectionName(collectionName))
                        .withFieldName("id")
                        .withIndexType(IndexType.AUTOINDEX)
                        .build());


        milvusServiceClient.createIndex(
                CreateIndexParam.newBuilder()
                        .withCollectionName(wrapCollectionName(collectionName))
                        .withFieldName("partition_name")
                        .withIndexType(IndexType.AUTOINDEX)
                        .build());

        milvusServiceClient.createIndex(
                CreateIndexParam.newBuilder()
                        .withCollectionName(wrapCollectionName(collectionName))
                        .withFieldName("content_vector_question")
                        .withIndexType(IndexType.IVF_FLAT)
                        .withMetricType(MetricType.COSINE)
                        .withSyncMode(Boolean.FALSE)
                        .withExtraParam("{\"nlist\":1024}")
                        .build());

        milvusServiceClient.createIndex(
                CreateIndexParam.newBuilder()
                        .withCollectionName(wrapCollectionName(collectionName))
                        .withFieldName("content_vector_answer")
                        .withIndexType(IndexType.IVF_FLAT)
                        .withMetricType(MetricType.COSINE)
                        .withSyncMode(Boolean.FALSE)
                        .withExtraParam("{\"nlist\":1024}")
                        .build());


    }

    public void loadCollection(String collectionName) {
        R<RpcStatus> response = milvusServiceClient.loadCollection(LoadCollectionParam.newBuilder()
                //集合名称
                .withCollectionName(wrapCollectionName(collectionName))
                .build());
    }

    public void releaseCollection(String collectionName) {
        R<RpcStatus> response = milvusServiceClient.releaseCollection(ReleaseCollectionParam.newBuilder()
                .withCollectionName(wrapCollectionName(collectionName))
                .build());
    }

    public List<String> searchContentByVector(String collectionName, String filterKey, String question, Integer topK) {
        List<List<Float>> questionVector = embeddingClient.doEmbedding(question);
        List<BaseMilvusDto> search = search(collectionName, questionVector, topK, filterKey);
        return search.stream().map(BaseMilvusDto::getContent).collect(Collectors.toList());
    }

    public List<BaseMilvusDto> search(String collectionName, String filterKey, String str, Integer topK) {
        List<List<Float>> string = embeddingClient.doEmbedding(str);
        return search(collectionName, string, topK, filterKey);
    }

    public List<BaseMilvusDto> search(String collectionName, List<List<Float>> questionVector, Integer topK,
                                      String filterKey) {
        loadCollection(collectionName);

        List<Long> lists = new ArrayList<>();
        List<Long> idsByQuestion = searchContent(collectionName, "content_vector_question", questionVector, topK, filterKey);
        List<Long> idsByAnswer = searchContent(collectionName, "content_vector_answer", questionVector, topK, filterKey);
        lists.addAll(idsByQuestion);
        lists.addAll(idsByAnswer);
        lists.removeIf(Objects::isNull);
        return getByIds(collectionName, lists);
    }

    private List<Long> searchContent(String collectionName, String vectorFieldName, List<List<Float>> vectors, Integer topK,
                                     String filterKey) {
        int nprobeVectorSize = vectors.get(0).size();
        String paramsInJson = "{\"nprobe\": " + nprobeVectorSize + "}";
        SearchParam.Builder contentVectorBuilder = SearchParam.newBuilder().withCollectionName(wrapCollectionName(collectionName))
                .withParams(paramsInJson)
                .withMetricType(MetricType.COSINE)
                .withVectors(vectors)
                .withVectorFieldName(vectorFieldName)
                .withTopK(topK);

        if (!ObjectUtil.isEmpty(filterKey)) {
            contentVectorBuilder.withExpr("partition_name == \"" + filterKey + "\"");
        }

        SearchParam searchParam = contentVectorBuilder.build();

        R<SearchResults> searchResultsR = milvusServiceClient.search(searchParam);
        SearchResults searchResultsRData = searchResultsR.getData();
        return searchResultsRData.getResults().getIds().getIntId().getDataList();
    }

    public List<BaseMilvusDto> getByIds(String collectionName, List<Long> ids) {
        if (ObjectUtil.isEmpty(ids)) {
            return new ArrayList<>();
        }
        GetIdsParam build = GetIdsParam.newBuilder().withCollectionName(wrapCollectionName(collectionName)).withPrimaryIds(ids)
                .withOutputFields(Lists.newArrayList("id", "partition_name", "content"))
                .build();
        R<GetResponse> response = milvusServiceClient.get(build);
        GetResponse getResponse = response.getData();
        List<Map<String, Object>> collect = getResponse.getRowRecords().stream().map(QueryResultsWrapper.RowRecord::getFieldValues).collect(Collectors.toList());

        return JSON.parseArray(JSON.toJSONString(collect), BaseMilvusDto.class);
    }


    public void doInsert(String collectionName, String filterName, List<QuestionAndAnswer> questionAnswers) {
        if (ObjectUtil.isEmpty(filterName)) {
            filterName = "default";
        }
        List<String> questions = questionAnswers.stream().map(QuestionAndAnswer::getQuestion).collect(Collectors.toList());
        List<List<Float>> questionsVector = embeddingClient.doEmbedding(questions);
        List<String> answers = questionAnswers.stream().map(QuestionAndAnswer::getAnswer).collect(Collectors.toList());
        List<List<Float>> answersVector = embeddingClient.doEmbedding(questions);
        if (questionsVector.size() != answers.size() || questionsVector.size() != questionAnswers.size()) {
            throw new BizRuntimeException("questionAnswers 数量不匹配");
        }
        List<InsertParam.Field> fields = new ArrayList<>();
        fields.add(new InsertParam.Field("content", answers));
        fields.add(new InsertParam.Field("partition_name", Lists.newArrayList(filterName)));
        fields.add(new InsertParam.Field("content_vector_question", questionsVector));
        fields.add(new InsertParam.Field("content_vector_answer", answersVector));
        InsertParam insertParam = InsertParam.newBuilder()
                .withCollectionName(wrapCollectionName(collectionName))
                .withFields(fields)
                .build();
        milvusServiceClient.insert(insertParam);
    }

    public void deleteEntity(String collectionName, String expr) {
        R<MutationResult> response = milvusServiceClient.delete(
                DeleteParam.newBuilder()
                        .withCollectionName(wrapCollectionName(collectionName))
                        //条件 如: id == 1
                        .withExpr(expr)
                        .build()
        );
    }

    public void deleteCollection(String collectionName) {
        DropCollectionParam build = DropCollectionParam.newBuilder().withCollectionName(wrapCollectionName(collectionName)).build();
        R<RpcStatus> rpcStatusR = milvusServiceClient.dropCollection(build);
    }
}
